package com.yoshino.leetcode.improve40.twentyth;

import java.util.*;

class Solution {

    public int[] topKFrequent(int[] nums, int k) {
        // 哈希 + 优先队列
        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            map.put(num, map.getOrDefault(num, 0) + 1);
        }
        PriorityQueue<int[]> pq = new PriorityQueue<>(Comparator.comparingInt(a -> a[0]));
        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            int num = entry.getValue();
            if (pq.size() == k) {
                // 需弹出
                if (pq.peek()[0] < num) {
                    pq.poll();
                    pq.offer(new int[]{num, entry.getKey()});
                }
            } else {
                pq.offer(new int[]{num, entry.getKey()});
            }
        }
        int[] res = new int[k];
        for (int i = 0; i < k; i++) {
            res[i] = pq.poll()[1];
        }
        return res;
    }

    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        // 优先队列
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (a, b) -> nums1[a[0]] + nums2[a[1]] - nums1[b[0]] - nums2[b[1]]);
        List<List<Integer>> res = new ArrayList<>();
        int m = nums1.length;
        int n = nums2.length;
        for (int i = 0; i < Math.min(m, k); i++) {
            pq.offer(new int[] {i, 0});
        }
        while (k-- > 0 && !pq.isEmpty()) {
            int[] poll = pq.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[poll[0]]);
            list.add(nums2[poll[1]]);
            res.add(list);
            if (poll[1] + 1 < n) {
                pq.offer(new int[] {poll[0], poll[1] + 1});
            }
        }
        return res;
    }
}